#include "stdafx.h"
#include <Psapi.h>
#include "Config.h"

#define IMAGE_VIRTUAL_TO_ABOSLUTE(type, hModule, virtualAdress) ((type)(DWORD(hModule) + DWORD(virtualAdress)))
#define IMPORT_DIRECTORY(hModule) IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_IMPORT_DESCRIPTOR, hModule, GetDataDirectory(hModule)[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress)
#define EXPORT_DIRECTORY(hModule) IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_EXPORT_DIRECTORY, hModule, GetDataDirectory(hModule)[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress)

PIMAGE_SECTION_HEADER GetFirstSection(LPVOID hModuleBase, LPWORD lpwNumberOfSections)
{
	try
	{
		PIMAGE_DOS_HEADER pimage_dos_header = PIMAGE_DOS_HEADER(hModuleBase);
		if (pimage_dos_header->e_magic != IMAGE_DOS_SIGNATURE)
		{
			Log("Uknown image format, DOS header not detected.");
			throw;
		}

		PIMAGE_NT_HEADERS pimage_nt_headers = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_NT_HEADERS, hModuleBase, pimage_dos_header->e_lfanew);
		if (pimage_nt_headers->Signature != IMAGE_NT_SIGNATURE)
		{
			Log("Uknown image format, NT header not detected.");
			throw;
		}

		*lpwNumberOfSections = pimage_nt_headers->FileHeader.NumberOfSections;

		return IMAGE_FIRST_SECTION(pimage_nt_headers);
	}
	catch (...)
	{
		Log("Error getting first section address. Image base: 0x%X", hModuleBase);
		throw;
	}
}

DWORD RvaToRawPointer(LPVOID hModuleBase, DWORD dwRVA, PIMAGE_SECTION_HEADER pFirstSection, WORD wNumberOfSections)
{
	try
	{
		for (int i = 0; i < wNumberOfSections; i++)
			if (dwRVA >= pFirstSection[i].VirtualAddress && dwRVA - pFirstSection[i].VirtualAddress < pFirstSection[i].SizeOfRawData)
				return (DWORD)hModuleBase + pFirstSection[i].PointerToRawData + (dwRVA - pFirstSection[i].VirtualAddress);

		throw;
	}
	catch (...)
	{
		Log("Error converting RVA to raw data pointer. Image base: 0x%X", hModuleBase);
		throw;
	}
}

PIMAGE_DATA_DIRECTORY GetDataDirectory(LPVOID hModuleBase, int iIndex)
{
	try
	{
		PIMAGE_DOS_HEADER pimage_dos_header = PIMAGE_DOS_HEADER(hModuleBase);
		if (pimage_dos_header->e_magic != IMAGE_DOS_SIGNATURE)
		{
			Log("Uknown image format, DOS header not detected.");
			throw;
		}

		PIMAGE_NT_HEADERS pimage_nt_headers = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_NT_HEADERS, hModuleBase, pimage_dos_header->e_lfanew);
		if (pimage_nt_headers->Signature != IMAGE_NT_SIGNATURE)
		{
			Log("Uknown image format, NT header not detected.");
			throw;
		}

		if (pimage_nt_headers->OptionalHeader.NumberOfRvaAndSizes <= iIndex)
		{
			Log("Directory of index %d not found, Number of supported directories is %d.", iIndex, pimage_nt_headers->OptionalHeader.NumberOfRvaAndSizes);
			throw;
		}

		return pimage_nt_headers->OptionalHeader.DataDirectory + iIndex;
	}
	catch (...)
	{
		Log("Error getting directory of index %d. Image base: 0x%X", iIndex, hModuleBase);
		throw;
	}
}

void WriteToProtectedMemory(LPDWORD lpTarget, DWORD value)
{
	try
	{
		DWORD dwOldValue;
		if (!VirtualProtect(lpTarget, sizeof(DWORD), PAGE_READWRITE, &dwOldValue))
		{
			LogWin32Error("Error: Couldn't change memory protection level to read/write.");
			throw;
		}
		*lpTarget = value;
		if (!VirtualProtect(lpTarget, sizeof(DWORD), dwOldValue, &dwOldValue))
		{
			LogWin32Error("[Error ignored] Couldn't change memory protection level back to original level. Location: 0x%X", lpTarget);
		}
	}
	catch (...)
	{
		Log("Error writing to location 0x%X. Value to be writter: 0x%X", lpTarget, value);
		throw;
	}
}

void HookImportTable(HMODULE hModule, LPCSTR strImportName, int iNumberOfProcs, LPCSTR astrProcNames[], LPVOID alpHookProc[], LPVOID alpOriginalProc[])
{
	try
	{
		char strBaseName[128];
		strcpy(strBaseName, GetModuleName(hModule));
		if (strBaseName)
			Log("Hooking import table for module [0x%X]: %s", hModule, strBaseName);
		else
			LogWin32Error("Hooking import table for module [0x%X]: <couldn't get name>", hModule);

		PIMAGE_DATA_DIRECTORY pImportDirectory = GetDataDirectory(hModule, IMAGE_DIRECTORY_ENTRY_IMPORT);

		if (pImportDirectory->Size == 0)
		{
			Log("Import directory size is zero, module is skipped");
			return;
		}

		BOOL bImportFound = FALSE;

		PIMAGE_IMPORT_DESCRIPTOR lpImports = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_IMPORT_DESCRIPTOR, hModule, pImportDirectory->VirtualAddress);

		while(lpImports->Name != 0)
		{
			LPCSTR strModuleName = IMAGE_VIRTUAL_TO_ABOSLUTE(LPCSTR, hModule, lpImports->Name);

			if (IsStringsEqualIgnoreCase(strImportName, strModuleName))
			{
				bImportFound = TRUE;
				break;
			}

			lpImports++;
		}

		if (!bImportFound)
		{
			Log("Library \"%s\" was not found in import tables, module skipped.", strImportName);
			return;
		}

		PIMAGE_THUNK_DATA pOriginalThunk;
		HANDLE hModuleFile;
		HANDLE hFileMapping;

		// Handle know bug with linkers not filling original think
		if (lpImports->OriginalFirstThunk == 0)
		{
			Log("Original thunk table is empty, attempting to load image file");

			char strFileName[1024];

			if (!GetModuleFileNameA(hModule, strFileName, sizeof(strFileName)))
			{
				LogWin32Error("Couldn't get image filename, skipping import");
				return;
			}

			hModuleFile = CreateFileA(strFileName, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, INVALID_HANDLE_VALUE);
			if (hModuleFile == INVALID_HANDLE_VALUE)
			{
				LogWin32Error("Couldn't open image file, skipping import");
				return;
			}

			hFileMapping = CreateFileMappingA(hModuleFile, NULL, PAGE_READONLY, 0, 0, NULL);
			if (hFileMapping == INVALID_HANDLE_VALUE)
			{
				LogWin32Error("Couldn't map image file into memory, skipping import");
				return;
			}

			LPVOID lpFileBase = MapViewOfFile(hFileMapping, FILE_MAP_READ, 0, 0, 0);
			if (lpFileBase == NULL)
			{
				LogWin32Error("Couldn't get map view over image file, skipping import");
				return;
			}

			WORD wNumberOfSections;
			PIMAGE_SECTION_HEADER pFirstSection = GetFirstSection(lpFileBase, &wNumberOfSections);
			pOriginalThunk = (PIMAGE_THUNK_DATA)RvaToRawPointer(lpFileBase, lpImports->FirstThunk, pFirstSection, wNumberOfSections);
		}
		else
		{
			pOriginalThunk = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_THUNK_DATA, hModule, lpImports->OriginalFirstThunk);
		}

		PIMAGE_THUNK_DATA pThunk = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_THUNK_DATA, hModule, lpImports->FirstThunk);

		BOOL bAtLeastOneMethodHooked = FALSE;

		while(pOriginalThunk->u1.AddressOfData != 0)
		{
			// Handle importing by ordinal
			if (IMAGE_SNAP_BY_ORDINAL32(pOriginalThunk->u1.Ordinal))
			{
				DWORD dwOrdinal = pOriginalThunk->u1.Ordinal & 0x7FFFFFFF;
				Log("Method imported by ordinal, skipped because it is not supported. Ordinal: %d", dwOrdinal);
			}
			else
			{
				PIMAGE_IMPORT_BY_NAME lpImportByName = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_IMPORT_BY_NAME, hModule, pOriginalThunk->u1.Function);
				LPCSTR lpFunctionName = (LPCSTR)lpImportByName->Name;
				for (int i = 0; i < iNumberOfProcs; i++)
					if (strcmp(astrProcNames[i], lpFunctionName) == 0)
					{
						Log("Hooking proc: %s", lpFunctionName);
						bAtLeastOneMethodHooked = TRUE;
						alpOriginalProc[i] = (LPVOID)pThunk->u1.AddressOfData;
						WriteToProtectedMemory((LPDWORD)pThunk, (DWORD)alpHookProc[i]);
					}
			}

			pOriginalThunk++;
			pThunk++;
		}

		if (!bAtLeastOneMethodHooked)
			Log("No matching proc names to hook for this library");

		// If image file was used, close file handles
		if (lpImports->OriginalFirstThunk == 0)
		{
			CloseHandle(hFileMapping);
			CloseHandle(hModuleFile);
		}
	}
	catch (...)
	{
		Log("Error hooking module: 0x%X for library: %s", hModule, strImportName);
	}
}

void Hook(LPCSTR strModule, HMODULE hCallingModule, int iNumberOfProcs, LPCSTR astrProcNames[], LPVOID alpHookProc[], LPVOID alpOriginalProc[])
{
	try
	{
		Log("=========================================================");
		Log("Hooking all import tables for library: %s", strModule);
		Log("=========================================================");

		HMODULE hModuleToBeHooked = GetModuleHandleA(strModule);

		if (hModuleToBeHooked == NULL)
		{
			LogWin32Error("Couldn't get module handle");
			return;
		}

		for (int i = 0; i < iNumberOfProcs; i++)
		{
			alpOriginalProc[i] = GetProcAddress(hModuleToBeHooked, astrProcNames[i]);
			if (alpOriginalProc[i] == NULL)
			{
				LogWin32Error("Couldn't get original proc address for proc \"%s\"", astrProcNames[i]);
				return;
			}
		}

		HMODULE modules[1024];
		DWORD cbNeeded;
		if (!EnumProcessModules(GetCurrentProcess(), modules, sizeof(modules), &cbNeeded))
		{
			LogWin32Error("Error enumurating loaded modules using PSAPI library");
			return;
		}
		DWORD dwModulesCount = cbNeeded / 4;

		LPVOID* dummy = new LPVOID[iNumberOfProcs];
		for (int i = 0; i < dwModulesCount; i++)
			if (modules[i] != hModuleToBeHooked && modules[i] != hCallingModule)
				HookImportTable(modules[i], strModule, iNumberOfProcs, astrProcNames, alpHookProc, dummy);
		delete[] dummy;
		Log("=========================================================");
	}
	catch (...)
	{
		Log("Error in hooking algorithm.");
	}
}